﻿//
// Created by QU on 24-9-28.
//

#include "server/server/Worker.h"

#include <iostream>

Worker::Worker(asio::io_context& io_context, std::stop_token worker_stop_token)
    : m_io_context(io_context),
      m_worker_work_guard(asio::make_work_guard(io_context)),
      m_worker_stop_token(std::move(worker_stop_token)),
      m_shared_session_stop_source(std::make_shared<std::stop_source>()),
      m_is_working(false)
{
    std::cout << "Worker::Worker() Creating worker" << std::endl;
}

Worker::~Worker()
{
    std::cout << "Worker::~Worker() Destroying worker" << std::endl;
    if (m_is_working.load())
    {
        forceStop();
    }
}

void Worker::start()
{
    m_is_working.store(true);
    m_worker_thread = std::jthread(
        [this]
        {
            try
            {
                while (!m_worker_stop_token.stop_requested())
                {
                    m_io_context.run_for(std::chrono::milliseconds(1000));
                }

                // 开始关闭过程
                std::cout << "Worker received token, starting shutdown process..." << std::endl;

                // 停止接受新的工作
                m_worker_work_guard.reset();

                // 请求停止所有会话
                m_shared_session_stop_source->request_stop();

                // 给予一些时间让会话优雅地关闭
                const auto shutdown_start = std::chrono::steady_clock::now();
                while (std::chrono::steady_clock::now() - shutdown_start < std::chrono::seconds(10))
                {
                    if (m_io_context.run_for(std::chrono::milliseconds(200)) == 0)
                    {
                        // 如果没有更多的异步操作，提前退出
                        break;
                    }
                }

                // 如果还有未完成的操作，强制关闭
                if (!m_io_context.stopped())
                {
                    std::cout << "Forcing io_context to stop..." << std::endl;
                    m_io_context.stop();
                }

                // 清理所有会话
                clearAllSessions();
            }
            catch (const std::exception& e)
            {
                std::cerr << "Error in worker thread: " << e.what() << std::endl;
            }

            m_is_working.store(false);
        });
}

void Worker::forceStop()
{
    std::cout << "Worker::forceStop(): Force stopping worker..." << std::endl;
    m_shared_session_stop_source->request_stop();
    m_worker_work_guard.reset();
    m_io_context.stop();
    if (m_worker_thread.joinable())
    {
        m_worker_thread.request_stop();
        m_worker_thread.join();
    }
    std::cout << "Worker::forceStop(): Clearing sessions..." << std::endl;
    std::scoped_lock lock(m_mutex);
    m_sessions_map.clear();
}

void Worker::addSession(const std::string& uuid, std::shared_ptr<Session> session)
{
    std::scoped_lock lock(m_mutex);
    m_sessions_map[uuid] = std::move(session);
}

void Worker::clearSession(const std::string& uuid)
{
    std::scoped_lock lock(m_mutex);
    if (const auto it = m_sessions_map.find(uuid); it != m_sessions_map.end())
    {
        auto session = it->second;
        m_sessions_map.erase(it);

        // 确保 session 已经关闭
        if (session->isStarted())
        {
            session->close();
        }
    }
}


void Worker::clearAllSessions()
{
    std::vector<std::shared_ptr<Session>> sessions_to_close;
    {
        std::scoped_lock lock(m_mutex);
        for (auto& val : m_sessions_map | std::views::values)
        {
            sessions_to_close.push_back(val);
        }
        m_sessions_map.clear();
    }

    // 在释放锁后关闭 sessions
    for (const auto& session : sessions_to_close)
    {
        if (session->isStarted())
        {
            session->close();
        }
    }
}


asio::io_context& Worker::getIOContext() const
{
    return m_io_context;
}

std::shared_ptr<std::stop_source> Worker::getSharedSessionStopSource() const
{
    return m_shared_session_stop_source;
}

std::optional<std::shared_ptr<Session>> Worker::getSession(const std::string& session_uuid)
{
    std::scoped_lock lock(m_mutex);
    const auto it = m_sessions_map.find(session_uuid);
    return (it != m_sessions_map.end()) ? std::make_optional(it->second) : std::nullopt;
}

size_t Worker::getSessionCount() const
{
    std::scoped_lock lock(m_mutex);
    return m_sessions_map.size();
}

std::vector<std::shared_ptr<Session>> Worker::getAllSessions() const
{
    std::scoped_lock lock(m_mutex);
    std::vector<std::shared_ptr<Session>> sessions;
    sessions.reserve(m_sessions_map.size());
    for (const auto& val : m_sessions_map | std::views::values)
    {
        sessions.push_back(val);
    }
    return sessions;
}

bool Worker::isWorking() const
{
    return m_is_working.load();
}
